#pragma once

#include "NvInfer.h"
#include <numeric>
#include <iostream>

using namespace std;

// 计算总共有多少个元素
inline int64_t volume(const nvinfer1::Dims& d) {
    return std::accumulate(d.d, d.d + d.nbDims, 1, 
            std::multiplies<int64_t>());
}

// 返回数据类型占用的字节大小
inline unsigned int getElementSize(nvinfer1::DataType t) {
    switch (t) {
        case nvinfer1::DataType::kINT32: return 4;
        case nvinfer1::DataType::kFLOAT: return 4;
        case nvinfer1::DataType::kHALF:  return 2;
        case nvinfer1::DataType::kBOOL:  return 1;
        case nvinfer1::DataType::kINT8:  return 1;
    }
    throw std::runtime_error("Invalid DataType.");
    return 0;
}

inline nvinfer1::DataType fieldTypeToDataType(
        const nvinfer1::PluginFieldType ftype) {
    switch (ftype) {
        case nvinfer1::PluginFieldType::kFLOAT32:
        {
            cout << "PluginFieldType is Float32" << std::endl;
            return nvinfer1::DataType::kFLOAT;
        }
        case nvinfer1::PluginFieldType::kFLOAT16:
        {
            cout << "PluginFieldType is Float16" << std::endl;
            return nvinfer1::DataType::kHALF;
        }
        case nvinfer1::PluginFieldType::kINT32:
        {
            cout << "PluginFieldType is Int32" << std::endl;
            return nvinfer1::DataType::kINT32;
        }
        case nvinfer1::PluginFieldType::kINT8:
        {
            cout << "PluginFieldType is Int8" << std::endl;
            return nvinfer1::DataType::kINT8;
        }
        default: 
            throw std::invalid_argument("No corresponding datatype for plugin field type");
    }
}